import codecs

discourse_sympol = {'<D%s>' % i for i in range(8)}


def clear_text(text):
    text = text.strip()
    for dis in discourse_sympol:
        text = text.replace(' ' + dis, '')
    text = text.split(' ')[:250]
    text = ' '.join(text)
    if not text.endswith('.'):
        last_index = text.rfind('.')
        if last_index != -1:
            text = text[:last_index + 1]
    return text


prompt_path = '/home/rickwwang/project_research/fairseq/data-bin/prompt.txt.test'
all_prompts = codecs.open(prompt_path, 'r', encoding='utf8').readlines()
all_prompts = map(lambda x: x.strip(), all_prompts)


# transformer-lm_bpe500_warmlr_format_pre2/27hyp.txt.1.topk40.t1.0
# transformer-lm_bpe500_warmlr_format2/40hyp.txt.1.topk40.t1.0
# bilstm_transformer-lm_bpe500_format_pre_srl_hier4_fix/25hyp.txt.own.prompt2srl2_500_wo_smooth_d1.1.topk10.t1.0
# bilstm_transformer-lm_bpe500_format_pre_srl_hier_pos_fix/17hyp.txt.own.prompt2srl2_500_wo_smooth_d1_verb_loss.1.topk10.t1.0
exp_file = []
# exp_file.append('fire_data/running/transformer-lm_bpe500_warmlr_format2/40hyp.txt.1.topk40.t1.0')
exp_file.append('fire_data/running/transformer-lm_bpe500_warmlr_format_pre2/27hyp.txt.1.topk20.t1.0')
# exp_file.append('fire_data/running/bilstm_transformer-lm_bpe500_format_pre_srl_hier4_fix/25hyp.txt.own.prompt2srl2_500_wo_smooth_d1.1.topk20.t1.0')
# exp_file.append('out/bilstm_transformer-lm_bpe500_format_pre_srl_hier_pos_fix/17hyp.txt.own.prompt2srl2_500_wo_smooth_d1_verb_loss.8.topk20.t1.0')
# exp_file.append('out/bilstm_transformer-lm_bpe500_format_pre_srl_hier_pos_fix/17hyp.txt.1.topk5.t1.0')
# exp_file.append('out/bilstm_transformer-lm_bpe500_format_pre_srl_hier_pos_fix/17tgt.txt')
# exp_file.append('fire_data/running/transformer-lm_bpe500_warmlr_format_pre2_memory_gate_gpt/37hyp.txt.1.topk20.t1.0')
# exp_file.append('fire_data/running/transformer-lm_bpe500_warmlr_format_pre2_coherence_gpt/30hyp.txt.1.topk20.t1.0')
# exp_file.append('fire_data/running/transformer-lm_bpe500_warmlr_format_pre2_memory_gate_coherence_gpt/49hyp.txt.1.topk20.t1.0')
# exp_file.append('out/transformer-lm_bpe500_warmlr_format_pre2_memory_gate_relevance_gpt_fix/8hyp.txt.1.topk20.t1.0')
# exp_file.append('out/transformer-lm_bpe500_warmlr_format_pre2_memory_gate_relevance_gpt_fix/26hyp.txt.1.topk20.t1.0')
# exp_file.append('fire_data/running/prompt2abs/14tgt.txt')
# exp_file.append('fire_data/running/prompt2abs/14hyp.txt.1.topk20.t1.0')
# exp_file.append('fire_data/running/prompt2abs/40hyp.txt.1.topk20.t1.0')
# exp_file.append('fire_data/running/promptdstory/13hyp.txt.1.topk20.t1.0')
# exp_file.append('fire_data/running/prompt2abs/14hyp.txt.1.topk40.t1.0')
# exp_file.append('fire_data/running/prompt2abs/14hyp.txt.1.topk80.t1.0')
# exp_file.append('fire_data/running/prompt2abs/14hyp.txt.1.topk20.t2.0')
# exp_file.append('fire_data/running/prompt2abs/14hyp.txt.1.topk20.t1.0.nr6')
exp_file.append('fire_data/running/transformer-lm_bpe500_warmlr_format_pre2_coref/26hyp.txt.1.topk20.t1.0')
exp_file.append('fire_data/running/promptdstory_discourse_two/13hyp.txt.1.min150.0.max300.topk20.t1.0')
exp_file.append('fire_data/running/promptdstory_discourse_two/29hyp.txt.1.min150.0.max300.topk20.t1.0')
exp_file.append('fire_data/running/promptdstory_discourse/8hyp.txt.1.min150.0.max300.topk20.t1.0')
exp_file.append('fire_data/running/promptdstory_discourse/30hyp.txt.1.min150.0.max300.topk20.t1.0')


base_path = '/home/rickwwang/project_research/'

for exp in exp_file:
    story_path = base_path + exp
    result_path = story_path + '.human.test'
    all_stories = codecs.open(story_path, 'r', encoding='utf8').readlines()

    promp_to_story = dict()
    for s in all_stories:
        s = s.strip()
        key = s.split(' <SEP> ')[0]
        value = s.split(' <SEP> ')[1]
        value = clear_text(value)
        if key in promp_to_story:
            promp_to_story[key].append(value)
        else:
            promp_to_story[key] = [value]

    format_stories = []
    for p in all_prompts:
        if p in promp_to_story:
            ss = promp_to_story[p]

            format_stories.append(p)
            format_stories.append(ss[0])
        else:
            format_stories.append('null')
            format_stories.append('null')

    with codecs.open(result_path, 'w', encoding='utf8') as fout:
        fout.write('\n'.join(format_stories))
        fout.write('\n')